import argparse
import json
import os
from tqdm import tqdm
from datasets import load_dataset, Dataset
from vllm import LLM, SamplingParams
from transformers import AutoTokenizer

argparser = argparse.ArgumentParser()
argparser.add_argument("--model", type=str, required=True)
argparser.add_argument("--tensor-parallel-size", type=int, default=1)
argparser.add_argument("--dataset", type=str, required=True)

args = argparser.parse_args()

llm = LLM(args.model, tensor_parallel_size=args.tensor_parallel_size)
tokenizer = AutoTokenizer.from_pretrained(args.model)

def prompt(data):
    prefix = data.get("prefix", "")
    suffix = data.get("fim_suffix", "")
    return {"prompt": f"<PRE>{prefix}<SUF>{suffix}<MID>"}

data = []

with open(args.dataset, 'r', encoding='utf-8') as f:
    for line in f.readlines():
        data.append(json.loads(line))

dataset = Dataset.from_list(data).map(prompt, num_proc=64)

prompts = dataset["prompt"]

outputs = llm.generate(prompts, SamplingParams(max_tokens=256))

results, answer = [], []

for i, output in tqdm(enumerate(outputs), total=len(outputs)):
    if output.prompt != dataset[i]["prompt"]:
        print("?")
    text = tokenizer.decode(output.outputs[0].token_ids, skip_special_tokens=True, clean_up_tokenization_spaces=True)
    results.append({
        "prefix": dataset[i]["prefix"],
        "fim_suffix": dataset[i]["fim_suffix"],
        "prompt": output.prompt,
        "middle": text
    })
    answer.append([text])

path, filename = args.dataset.rsplit("/", 1)
filename = filename.split(".")[0]

with open(os.path.join(path, f"{filename}-output.jsonl"), 'w', encoding='utf-8') as f:
    for entry in results:
        f.write(json.dumps(entry, ensure_ascii=False)+"\n")

with open(os.path.join(path, f"{filename}-result.json"), 'w', encoding='utf-8') as f:
    json.dump(answer, f, ensure_ascii=False, indent=4)